import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchvision
import numpy as np
import matplotlib.pyplot as plt
# from torch.onnx.symbolic_opset9 import view
from torchvision import datasets, transforms

#对数据做归一化（-1， 1）
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize(0.5,0.5)
])

train_ds = torchvision.datasets.MNIST('data',train=True,transform=transform,download=False)
dataloader = torch.utils.data.DataLoader(train_ds,batch_size=64,shuffle=True)
imgs,_ =next(iter(dataloader))

class Generator(nn.Module):
    def __init__(self):
        super(Generator,self).__init__()
        self.main = nn.Sequential(
                                nn.Linear(100,256),
                                nn.ReLU(),
                                nn.Linear(256,512),
                                nn.ReLU(),
                                nn.Linear(512,28*28),
                                nn.Tanh() #(-1,1之间)
        )
    def forward(self,x):  #x表示长度为100的噪声输入
        img = self.main(x)
        img = img.view(-1,1,28,28)#变成28*28*1的图片
        return img

class Discriminator(nn.Module):
    def __init__(self):
        super(Discriminator,self).__init__()
        self.main = nn.Sequential(
                                nn.Linear(28*28,512),
                                nn.LeakyReLU(),
                                nn.Linear(512,256),
                                nn.LeakyReLU(),
                                nn.Linear(256,1),
                                nn.Sigmoid()
        )
    def forward(self,x): #x表示28*28的图片
        x = x.view(-1,28*28)#展平
        x = self.main(x)
        return x

device = 'mps' if torch.backends.mps.is_available() else 'cpu'
# device = 'cuda' if torch.cuda.is_available() else 'cpu'
print(device)
gen = Generator().to(device)
dis = Discriminator().to(device)
d_optim = torch.optim.Adam(dis.parameters(),lr=0.0001)
g_optim = torch.optim.Adam(gen.parameters(),lr=0.0001)
loss_fn = torch.nn.BCELoss()
def gen_img_plot(model,test_input):
    prediction = np.squeeze(model(test_input).detach().cpu().numpy())
    fig = plt.figure(figsize=(4,4))
#     for i in range(prediction.size(0)):
    for i in range(16):
        plt.subplot(4,4,i+1)
        plt.imshow((prediction[i]+1)/2)
        plt.axis('off')
    plt.show()

test_input = torch.randn(16,100,device=device)

D_loss = []
G_loss = []

from torch import detach

# 训练循环
for epoch in range(20):
    d_epoch_loss = 0
    g_epoch_loss = 0
    count = len(dataloader)
    for step, (img, _) in enumerate(dataloader):
        img = img.to(device)
        size = img.size(0)
        random_noise = torch.randn(size, 100, device=device)

        # 判别器训练
        d_optim.zero_grad()
        real_output = dis(img)  # 判别器输入真实的图片，real_output对真实图片的预测结果
        #         real_output = torch.sigmoid(real_output)
        d_real_loss = loss_fn(real_output, torch.ones_like(real_output))  # 判别器在真实图像上的损失
        d_real_loss.backward()

        # 判别器计算生成图像计的损失
        gen_img = gen(random_noise)
        fake_output = dis(gen_img.detach())  # 判别器输入生成的腿片，fake_output是对生成图片的预测结果
        d_fake_loss = loss_fn(fake_output, torch.zeros_like(fake_output))  # 判别器在生成图像上的损失
        d_fake_loss.backward()

        d_loss = d_real_loss + d_fake_loss
        d_optim.step()

        g_optim.zero_grad()
        fake_output = dis(gen_img)
        g_loss = loss_fn(fake_output, torch.ones_like(fake_output))  # 生成器的损失
        g_loss.backward()
        g_optim.step()

        with torch.no_grad():
            d_epoch_loss += d_loss
            g_epoch_loss += g_loss

    with torch.no_grad():
        d_epoch_loss /= count
        g_epoch_loss /= count
        D_loss.append(d_epoch_loss)
        G_loss.append(g_epoch_loss)
        print('Epoch:', epoch)
        gen_img_plot(gen, test_input)